from typing import TypeVar

import numpy as np

from zkl_aiutils_datasets.basics import Dataset
from zkl_aiutils_datasets.processing import PickedDataset

AnySample = TypeVar('AnySample')


class PickingShuffledDataset(PickedDataset[AnySample]):
    def __init__(self, dataset: Dataset[AnySample], *, random_seed: int = 42):
        indices = np.arange(len(dataset))
        np.random.default_rng(random_seed).shuffle(indices)
        indices = indices.tolist()
        super().__init__(dataset, indices)
        self._random_seed = random_seed

    @property
    def random_seed(self) -> int:
        return self._random_seed
